package cn.wzl.recognition.mvc;

import cn.wzl.recognition.bpnn.PictureRecognitionBPNN;
import cn.wzl.recognition.bpnn.util.PictureRecognitionConverter;
import cn.wzl.recognition.training.Trainer;

import javax.imageio.ImageIO;
import javax.imageio.stream.FileImageOutputStream;
import javax.imageio.stream.ImageOutputStream;
import java.awt.AWTException;
import java.awt.BasicStroke;
import java.awt.Color;
import java.awt.Graphics;
import java.awt.Graphics2D;
import java.awt.Rectangle;
import java.awt.Robot;
import java.awt.event.MouseEvent;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.Optional;

public class Model {

    private final View view;
    private final boolean isDebugMode;
    private final PictureRecognitionBPNN bp;
    private final Trainer trainer;

    private int startX;
    private int startY;

    public Model(View view, PictureRecognitionBPNN bp, Trainer trainer, boolean isDebugMode) {
        this.view = view;
        this.isDebugMode = isDebugMode;
        this.trainer = trainer;
        this.bp = bp;
    }

    public void mouseDragged(MouseEvent e) {
        Graphics2D graphics = (Graphics2D) view.getPicPanel().getGraphics();
        int endX = e.getX() - view.getInsets().left;
        int endY = e.getY() - view.getInsets().top;
        graphics.setColor(Color.black);
        graphics.setStroke(new BasicStroke(view.PEN_WIDTH));
        graphics.drawLine(startX, startY, endX, endY);
        startX = endX;
        startY = endY;
    }

    public void mousePressed(MouseEvent e) {
        startX = e.getX() - view.getInsets().left;
        startY = e.getY() - view.getInsets().top;
    }

    public void clearPanel() {
        Graphics graphics = view.getPicPanel().getGraphics();
        graphics.clearRect(0, 0, view.getPicPanel().getSize().width, view.getPicPanel().getSize().height);
    }

    public View getView() {
        return view;
    }

    public void run() {
        if (isDebugMode) {
            captureImageToFile("output/test/capture.jpg");
        }
        int ret = bp.run(getImage());

        if (isDebugMode) {
            if(ret < 0) {
                System.out.println("can not recognize the number");
            }
            else {
                System.out.println("this is " + ret);
            }
        }
    }

    public void trainingCurrentPic(int expectedValue, int times) {
        double [] image = getImage();
        while(times -- != 0) {
            bp.training(image, PictureRecognitionConverter.convertNumberToOutputArray(expectedValue));
        }
        System.out.println("training end.");
    }

    private double[] getImage() {
        return captureScreenInView()
                .map(this::convertToLittlePic)
                .map(this::retrievePicGrayArr)
                .orElseThrow(RuntimeException::new);
    }

    private Optional<BufferedImage> captureScreenInView() {
        try {
            return Optional.of(
                    new Robot().createScreenCapture(
                            new Rectangle(
                                    view.getX() + view.getInsets().left, view.getY() + view.getInsets().top,
                                    view.getPicPanel().getSize().width, view.getPicPanel().getSize().height
                            )
                    )
            );
        } catch (AWTException e) {
            e.printStackTrace();
            return Optional.empty();
        }
    }

    private BufferedImage convertToLittlePic(BufferedImage bufferedImage) {
        BufferedImage littlePic = new BufferedImage(28, 28, BufferedImage.TYPE_INT_RGB);
        littlePic.getGraphics().drawImage(bufferedImage, 0, 0, 28, 28, null);
        return littlePic;
    }

    private double[] retrievePicGrayArr(BufferedImage image) {
        double[] bitPic = new double[image.getWidth() * image.getHeight()];

        for (int i = 0; i < image.getWidth(); i++) {
            for (int j = 0; j < image.getHeight(); j++) {
                int rgb = image.getRGB(j, i);
                int gray = (
                        ((rgb & 0xff0000) >> 16) +
                                ((rgb & 0xff00) >> 8) +
                                (rgb & 0xff)
                ) / 3;
                int bitValue = (gray >= 255 / 2) ? 0 : 1;
                bitPic[i * image.getHeight() + j] = bitValue;
            }
        }
        if (isDebugMode) {
            printInputValue(bitPic, image.getWidth(), image.getHeight());
            bitMapOutput("output/test/capture_small.jpg", bitPic, image.getWidth(), image.getHeight());
        }
        return bitPic;
    }

    private void bitMapOutput(@SuppressWarnings("SameParameterValue") String fileName, double[] bitPic, int width, int height) {
        try {
            ImageOutputStream imageOutput = new FileImageOutputStream(new File(fileName));
            BufferedImage bufferedImage = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
            Graphics graphics = bufferedImage.getGraphics();
            for (int i = 0; i < width; i++) {
                for (int j = 0; j < height; j++) {
                    if(bitPic[i * height + j] > 0.5) {
                        graphics.setColor(Color.BLACK);
                        graphics.drawRect(j, i, 1, 1);
                    }
                    else{
                        graphics.setColor(Color.WHITE);
                        graphics.drawRect(j, i, 1, 1);
                    }
                }
            }
            ImageIO.write(bufferedImage, "jpg", imageOutput);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private void printInputValue(double[] pic, int width, int height) {
        for (int i = 0; i < width; i++) {
            for (int j = 0; j < height; j++) {
                System.out.print(pic[i * height + j]);
                System.out.print(' ');
            }
            System.out.print("\n");
        }
    }

    private void captureImageToFile(@SuppressWarnings("SameParameterValue") String fileName) {
        captureScreenInView().ifPresent(image -> {
            try {
                ImageIO.write(image, "JPEG", new FileOutputStream(fileName));
            } catch (IOException e) {
                e.printStackTrace();
            }
        });
    }

    public void training(int times) {
        this.trainer.training(times);
    }

}
